# -*- coding: utf-8 -*-

"""
边框颜色分析工具
用于分析PDF页面图像中的边框颜色
"""

import os
import argparse
from PIL import Image
import numpy as np
from collections import Counter

class BorderColorAnalyzer:
    def __init__(self):
        self.color_threshold = 200  # 用于识别边缘的阈值
        self.edge_sample_size = 1000  # 边缘采样点数量
    
    def load_image(self, image_path):
        """加载图像文件"""
        try:
            img = Image.open(image_path)
            if img.mode != 'RGB':
                img = img.convert('RGB')
            return img
        except Exception as e:
            print(f"加载图像失败: {str(e)}")
            return None
    
    def detect_edges(self, img):
        """使用简单的边缘检测方法"""
        width, height = img.size
        pixels = np.array(img)
        
        # 创建灰度图像
        gray = np.mean(pixels, axis=2).astype(np.uint8)
        
        # 简单边缘检测 (计算梯度)
        edges = np.zeros_like(gray)
        
        # 计算水平方向梯度
        for i in range(width - 1):
            for j in range(height):
                edges[j, i] = abs(int(gray[j, i]) - int(gray[j, i+1]))
        
        # 计算垂直方向梯度
        vertical_edges = np.zeros_like(gray)
        for i in range(width):
            for j in range(height - 1):
                vertical_edges[j, i] = abs(int(gray[j, i]) - int(gray[j+1, i]))
        
        # 合并边缘
        edges = np.maximum(edges, vertical_edges)
        
        # 二值化
        edges[edges < self.color_threshold] = 0
        edges[edges >= self.color_threshold] = 1
        
        return edges
    
    def collect_edge_colors(self, img, edges):
        """收集边缘点的颜色"""
        width, height = img.size
        pixels = np.array(img)
        
        # 找出所有边缘点
        edge_points = np.argwhere(edges > 0)
        
        if len(edge_points) == 0:
            print("未检测到边缘点")
            return []
        
        # 随机采样
        sample_size = min(self.edge_sample_size, len(edge_points))
        sample_indices = np.random.choice(len(edge_points), sample_size, replace=False)
        sample_points = edge_points[sample_indices]
        
        # 收集颜色
        edge_colors = []
        for j, i in sample_points:
            if 0 <= i < width and 0 <= j < height:
                color = tuple(pixels[j, i])
                edge_colors.append(color)
        
        return edge_colors
    
    def find_dominant_colors(self, colors, top_n=5):
        """找出主导颜色"""
        if not colors:
            return []
        
        # 统计颜色出现频率
        color_counter = Counter(colors)
        
        # 获取出现频率最高的颜色
        dominant_colors = color_counter.most_common(top_n)
        
        return dominant_colors
    
    def analyze_image(self, image_path):
        """分析图像中的边框颜色"""
        # 加载图像
        img = self.load_image(image_path)
        if img is None:
            return None
        
        print(f"成功加载图像: {image_path}")
        print(f"图像尺寸: {img.size}")
        
        # 检测边缘
        edges = self.detect_edges(img)
        
        # 统计边缘像素比例
        edge_pixel_count = np.sum(edges)
        total_pixels = edges.size
        edge_ratio = edge_pixel_count / total_pixels
        print(f"边缘像素比例: {edge_ratio:.4f} ({edge_pixel_count} / {total_pixels})")
        
        # 收集边缘颜色
        edge_colors = self.collect_edge_colors(img, edges)
        print(f"收集到 {len(edge_colors)} 个边缘颜色样本")
        
        # 找出主导颜色
        dominant_colors = self.find_dominant_colors(edge_colors)
        
        # 输出结果
        print("主导颜色 (RGB值, 出现次数):")
        for color, count in dominant_colors:
            print(f"{color}: {count}")
        
        return dominant_colors
    
    def analyze_directory(self, directory_path):
        """分析目录中的所有图像"""
        if not os.path.exists(directory_path):
            print(f"目录不存在: {directory_path}")
            return
        
        # 获取目录中的所有图像文件
        image_extensions = ['.png', '.jpg', '.jpeg', '.bmp']
        image_files = []
        
        for file in os.listdir(directory_path):
            if any(file.lower().endswith(ext) for ext in image_extensions):
                image_files.append(os.path.join(directory_path, file))
        
        print(f"在目录中找到 {len(image_files)} 个图像文件")
        
        # 分析每个图像
        all_dominant_colors = []
        
        for image_path in image_files:
            print(f"\n分析图像: {os.path.basename(image_path)}")
            dominant_colors = self.analyze_image(image_path)
            if dominant_colors:
                all_dominant_colors.extend(dominant_colors)
        
        # 汇总所有图像的主导颜色
        if all_dominant_colors:
            print("\n所有图像的主导颜色汇总:")
            # 重新统计
            color_counter = Counter([color for color, _ in all_dominant_colors])
            overall_dominant = color_counter.most_common(5)
            
            for color, count in overall_dominant:
                print(f"{color}: {count}")
    
    def analyze_rectangle_files(self, directory_path):
        """专门分析矩形文件中的边框颜色"""
        if not os.path.exists(directory_path):
            print(f"目录不存在: {directory_path}")
            return
        
        # 获取目录中的矩形图像文件
        rectangle_files = []
        
        for file in os.listdir(directory_path):
            if 'rectangle' in file.lower() and file.lower().endswith('.png'):
                rectangle_files.append(os.path.join(directory_path, file))
        
        print(f"在目录中找到 {len(rectangle_files)} 个矩形图像文件")
        
        # 分析每个矩形图像
        all_border_colors = []
        
        for image_path in rectangle_files:
            print(f"\n分析矩形图像: {os.path.basename(image_path)}")
            
            # 加载图像
            img = self.load_image(image_path)
            if img is None:
                continue
            
            width, height = img.size
            pixels = np.array(img)
            
            # 只采样边缘像素 (边框)
            border_colors = []
            border_width = 5  # 边框宽度采样范围
            
            # 上边框
            for i in range(width):
                for j in range(min(border_width, height)):
                    color = tuple(pixels[j, i])
                    border_colors.append(color)
            
            # 下边框
            for i in range(width):
                for j in range(max(0, height - border_width), height):
                    color = tuple(pixels[j, i])
                    border_colors.append(color)
            
            # 左边框
            for i in range(min(border_width, width)):
                for j in range(height):
                    color = tuple(pixels[j, i])
                    border_colors.append(color)
            
            # 右边框
            for i in range(max(0, width - border_width), width):
                for j in range(height):
                    color = tuple(pixels[j, i])
                    border_colors.append(color)
            
            # 找出边框的主导颜色
            if border_colors:
                dominant_border = self.find_dominant_colors(border_colors, top_n=3)
                print(f"边框主导颜色: {dominant_border}")
                all_border_colors.extend(border_colors)
        
        # 汇总所有矩形的边框颜色
        if all_border_colors:
            print("\n所有矩形的边框颜色汇总:")
            color_counter = Counter(all_border_colors)
            overall_dominant = color_counter.most_common(5)
            
            for color, count in overall_dominant:
                print(f"{color}: {count}")


def main():
    parser = argparse.ArgumentParser(description='边框颜色分析工具')
    parser.add_argument('--image-path', type=str, help='图像文件路径')
    parser.add_argument('--directory-path', type=str, help='图像目录路径')
    parser.add_argument('--analyze-rectangles', type=str, help='分析矩形文件的边框颜色')
    args = parser.parse_args()
    
    analyzer = BorderColorAnalyzer()
    
    if args.image_path:
        analyzer.analyze_image(args.image_path)
    elif args.directory_path:
        analyzer.analyze_directory(args.directory_path)
    elif args.analyze_rectangles:
        analyzer.analyze_rectangle_files(args.analyze_rectangles)
    else:
        print("请指定图像文件路径或目录路径")


if __name__ == '__main__':
    main()